package misc;

import java.util.AbstractCollection;
import java.util.AbstractSet;
import java.util.Collection;
import java.util.IdentityHashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicReference;

public class ConcurrentIdentityHashMap<K, V> implements ConcurrentMap<K, V> {

    private final AtomicReference<IdentityHashMap<K, V>> delegateHolder = new AtomicReference<IdentityHashMap<K, V>>(new IdentityHashMap<K, V>());

    public V putIfAbsent(K key, V value) {
        checkNull(key);
        checkNull(value);
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            V oldValue = current.get(key);
            if (oldValue != null) {
                return oldValue;
            }
            IdentityHashMap<K, V> copy = copy(current);
            copy.put(key, value);
            if (delegateHolder.compareAndSet(current, copy)) {
                return null;
            }
        }
    }

    @SuppressWarnings("unchecked")
    private IdentityHashMap<K, V> copy(IdentityHashMap<K, V> current) {
        return (IdentityHashMap<K, V>) current.clone();
    }

    private void checkNull(Object any) {
        if (any == null) {
            throw new NullPointerException("null");
        }
    }

    public boolean remove(Object key, Object value) {
        checkNull(key);
        checkNull(value);
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            V oldValue = current.get(key);
            if (oldValue == null || !oldValue.equals(value)) {
                return false;
            }
            IdentityHashMap<K, V> copy = copy(current);
            copy.remove(key);
            if (delegateHolder.compareAndSet(current, copy)) {
                return true;
            }
        }
    }

    public V replace(K key, V value) {
        checkNull(key);
        checkNull(value);
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            if (!current.containsKey(key)) {
                return null;
            }
            IdentityHashMap<K, V> copy = copy(current);
            V oldValue = copy.put(key, value);
            if (delegateHolder.compareAndSet(current, copy)) {
                return oldValue;
            }
        }
    }

    public boolean replace(K key, V oldValue, V newValue) {
        checkNull(key);
        checkNull(oldValue);
        checkNull(newValue);
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            V oldValue2 = current.get(key);
            if (oldValue == null || !oldValue2.equals(oldValue)) {
                return false;
            }
            IdentityHashMap<K, V> copy = copy(current);
            copy.remove(key);
            if (delegateHolder.compareAndSet(current, copy)) {
                return true;
            }
        }
    }

    public void clear() {
        delegateHolder.set(new IdentityHashMap<K, V>());
    }

    public boolean containsKey(Object key) {
        return delegateHolder.get().containsKey(key);
    }

    public boolean containsValue(Object value) {
        return delegateHolder.get().containsValue(value);
    }

    public V get(Object key) {
        return delegateHolder.get().get(key);
    }

    public boolean isEmpty() {
        return delegateHolder.get().isEmpty();
    }

    public V put(K key, V value) {
        checkNull(key);
        checkNull(value);
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            IdentityHashMap<K, V> copy = copy(current);
            V old = copy.put(key, value);
            if (delegateHolder.compareAndSet(current, copy)) {
                return old;
            }
        }
    }

    public void putAll(Map<? extends K, ? extends V> m) {
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            IdentityHashMap<K, V> copy = copy(current);
            copy.putAll(m);
            if (delegateHolder.compareAndSet(current, copy)) {
                return;
            }
        }
    }

    public V remove(Object key) {
        while (true) {
            IdentityHashMap<K, V> current = delegateHolder.get();
            IdentityHashMap<K, V> copy = copy(current);
            V old = copy.remove(key);
            if (delegateHolder.compareAndSet(current, copy)) {
                return old;
            }
        }
    }

    public int size() {
        return delegateHolder.get().size();
    }

    public Collection<V> values() {
        return new ValuesCollection();
    }

    public Set<K> keySet() {
        return new KeySet();
    }

    public Set<Map.Entry<K, V>> entrySet() {
        return new EntrySet();
    }

    private final class EntrySet extends AbstractSet<Entry<K, V>> {
        private final class EntrySetIterator implements Iterator<Entry<K, V>> {
            private final class ForwardingEntry implements Entry<K, V> {
                private final Entry<K, V> localNext;

                private ForwardingEntry(Entry<K, V> localNext) {
                    this.localNext = localNext;
                }

                public K getKey() {
                    return localNext.getKey();
                }

                public V getValue() {
                    return localNext.getValue();
                }

                public V setValue(V value) {
                    return ConcurrentIdentityHashMap.this.put(lastReturned.getKey(), value);
                }
            }

            final Iterator<Entry<K, V>> delegate = map.entrySet().iterator();
            volatile Entry<K, V> lastReturned;

            public boolean hasNext() {
                return delegate.hasNext();
            }

            public Entry<K, V> next() {
                final Entry<K, V> localNext = delegate.next();
                lastReturned = localNext;
                return new ForwardingEntry(localNext);
            }

            public void remove() {
                ConcurrentIdentityHashMap.this.remove(lastReturned.getKey());
            }
        }

        private final Map<K, V> map = delegateHolder.get();

        @Override
        public Iterator<Entry<K, V>> iterator() {
            return new EntrySetIterator();
        }

        @Override
        public int size() {
            return map.size();
        }
    }

    private final class KeySet extends AbstractSet<K> {
        private final class KeysIterator implements Iterator<K> {
            final Iterator<Entry<K, V>> delegate = map.entrySet().iterator();
            volatile Entry<K, V> lastReturned;

            public boolean hasNext() {
                return delegate.hasNext();
            }

            public K next() {
                final Entry<K, V> localNext = delegate.next();
                lastReturned = localNext;
                return localNext.getKey();
            }

            public void remove() {
                ConcurrentIdentityHashMap.this.remove(lastReturned.getKey());
            }
        }

        private final Map<K, V> map = delegateHolder.get();

        @Override
        public Iterator<K> iterator() {
            return new KeysIterator();
        }

        @Override
        public int size() {
            return map.size();
        }
    }

    private final class ValuesCollection extends AbstractCollection<V> {
        private final class ValuesIterator implements Iterator<V> {
            final Iterator<Entry<K, V>> delegate = map.entrySet().iterator();
            volatile Entry<K, V> lastReturned;

            public boolean hasNext() {
                return delegate.hasNext();
            }

            public V next() {
                final Entry<K, V> localNext = delegate.next();
                lastReturned = localNext;
                return localNext.getValue();
            }

            public void remove() {
                ConcurrentIdentityHashMap.this.remove(lastReturned.getKey());
            }
        }

        private final Map<K, V> map = delegateHolder.get();

        @Override
        public Iterator<V> iterator() {
            return new ValuesIterator();
        }

        @Override
        public int size() {
            return map.size();
        }
    }

}
